# -*- coding: utf-8 -*-

import sys
from PyQt5 import QtCore, QtGui
from PyQt5.QtWidgets import QApplication, QWidget, QPushButton, QTextEdit, QGridLayout, QLabel
from PyQt5.QtGui import QFont, QPixmap, QImage, QPainter, QPen
from PyQt5.QtCore import Qt, QEvent
from flow.label_video.label_video_backend import LabelVideoBackend


class LabelVideoWindow(QWidget):
    def __init__(self, root_folder, video_folder):
        super(LabelVideoWindow, self).__init__()
        self.status_line_edit = QTextEdit()
        self.previous_image_button = QPushButton("返回上一张(A)")
        self.next_image_button = QPushButton("进入下一张(D)")
        self.image_label = QLabel()
        self.backend = LabelVideoBackend(root_folder, video_folder)

        # used for image label action
        self.current_image = None
        self.top_left_position = None

        self.previous_image_button.clicked.connect(self.previous_image_button_pressed)
        self.next_image_button.clicked.connect(self.next_image_button_pressed)

        self.init_ui()
        self.set_up_child_widgets_property()
        self.setGeometry(300, 300, 1000, 550)
        self.show_image()

    def init_ui(self):
        font = QFont("宋体", 15)
        font.setBold(True)
        self.status_line_edit.setFixedHeight(40)
        self.previous_image_button.setFixedHeight(40)
        self.previous_image_button.setFont(font)
        self.next_image_button.setFixedHeight(40)
        self.next_image_button.setFont(font)
        grid_layout = QGridLayout()
        grid_layout.addWidget(self.status_line_edit, 0, 0, 1, 8)
        grid_layout.addWidget(self.previous_image_button, 0, 8, 1, 2)
        grid_layout.addWidget(self.next_image_button, 0, 10, 1, 2)
        grid_layout.setSpacing(15)
        grid_layout.addWidget(self.image_label, 2, 0, 2, 12)
        self.setLayout(grid_layout)
        self.setWindowTitle("客流标注工具")

    def set_up_child_widgets_property(self):
        self.status_line_edit.setDisabled(True)
        self.image_label.setScaledContents(True)
        self.image_label.installEventFilter(self)
        self.image_label.setMouseTracking(True)
        self.image_label.setFocus()

    def show_label_status(self, processed_count, left_count):
        text = "<span style=\"font-size:18pt; font-weight:600; color:#ff0000;\">"
        text += "正在标注。。。"
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += "已经标注了 "
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#0000ff;\">"
        text += "%d" % processed_count
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += " 张图片"
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += ", 剩余 "
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#0000ff;\">"
        text += "%d" % left_count
        text += "</span>"
        text += "<span style=\"font-size:18pt; font-weight:600; color:#000000;\">"
        text += " 张图片"
        text += "</span>"
        self.status_line_edit.setText(text)

    def show_image(self):
        image = self.backend.load_image()
        height, width, channel = image.shape
        q_image = QImage(image.data, width, height, 3*width, QImage.Format_RGB888).rgbSwapped()
        self.current_image = q_image
        self.image_label.setPixmap(QPixmap(q_image))

        labeled_frame_number = self.backend.labeled_frame_number
        labeled_frame_index = self.backend.labeled_frame_index
        if labeled_frame_index == 0:
            self.previous_image_button.setDisabled(True)
        else:
            self.previous_image_button.setDisabled(False)
        if labeled_frame_index == labeled_frame_number - 1:
            self.next_image_button.setDisabled(True)
        else:
            self.next_image_button.setDisabled(False)
        self.show_label_status(labeled_frame_index + 1,
                               labeled_frame_number - labeled_frame_index - 1)

        super(LabelVideoWindow, self).show()

    def draw_dashed_lines(self, event):
        image = self.current_image.copy()
        x = event.x() / self.image_label.width() * image.width()
        y = event.y() / self.image_label.height() * image.height()
        painter = QPainter()
        painter.begin(image)
        painter.setPen(QPen(QtCore.Qt.yellow, 2, QtCore.Qt.DashDotLine))
        painter.drawLine(0, y, self.image_label.width() - 1, y)
        painter.drawLine(x, 0, x, self.image_label.height() - 1)
        if self.top_left_position is not None:
            painter.setPen(QPen(QPen(QtCore.Qt.blue, 2)))
            x1, y1 = self.top_left_position
            painter.drawRect(x1, y1, x - x1, y - y1)
        painter.end()
        self.image_label.setPixmap(QPixmap(image))

    def draw_rect(self, event):
        if self.top_left_position is None:
            return
        x1, y1 = self.top_left_position
        x2 = event.x() / self.image_label.width() * self.current_image.width()
        y2 = event.y() / self.image_label.height() * self.current_image.height()
        if x2 - x1 < 10 or y2 - y1 < 10:
            return
        image = self.current_image.copy()
        painter = QPainter()
        painter.begin(image)
        painter.setPen(QPen(QPen(QtCore.Qt.red, 2)))
        painter.drawRect(x1, y1, x2 - x1 + 1, y2 - y1 + 1)
        painter.end()
        self.current_image = image.copy()
        self.image_label.setPixmap(QPixmap(image))

    def eventFilter(self, source, event):
        if source is not self.image_label:
            return QWidget.eventFilter(self, source, event)
        if event.type() == QEvent.MouseMove:
            self.draw_dashed_lines(event)
        if event.type() == QEvent.MouseButtonPress and event.button() == Qt.LeftButton:
            self.top_left_position = (
                event.x() / self.image_label.width() * self.current_image.width(),
                event.y() / self.image_label.height() * self.current_image.height()
            )
        if event.type() == QEvent.MouseButtonRelease and event.button() == Qt.LeftButton:
            self.draw_rect(event)
            x1 = self.top_left_position[0] / self.current_image.width()
            y1 = self.top_left_position[1] / self.current_image.height()
            x2 = event.x() / self.image_label.width()
            y2 = event.y() / self.image_label.height()
            self.backend.add_one_rectangle([x1, y1, x2, y2])
            self.show_image()
            self.top_left_position = None
        if event.type() == QEvent.MouseButtonRelease and event.button() == Qt.RightButton:
            self.backend.remove_one_rectangle()
            self.show_image()
        return QWidget.eventFilter(self, source, event)

    def previous_image_button_pressed(self):
        self.backend.labeled_frame_index -= 1
        self.show_image()

    def next_image_button_pressed(self):
        self.backend.labeled_frame_index += 1
        self.show_image()

    def event(self, event):
        if isinstance(event, QtGui.QKeyEvent) and event.type() == QtCore.QEvent.KeyPress:
            if event.key() == Qt.Key_A:
                self.previous_image_button_pressed()
            elif event.key() == Qt.Key_D:
                self.next_image_button_pressed()
            else:
                pass
        return super().event(event)

    def closeEvent(self, a0: QtGui.QCloseEvent):
        self.backend.save_labeled_result()


def main():
    root_folder = r"F:\data\bus_videos\59092_00E0B453CE0B\2019_3_18"
    video_path = r"F:\data\bus_videos\59092_00E0B453CE0B\2019_3_18\c0\2019_3_18-7_5_7-0(9_0_9_0).avi"
    app = QApplication(sys.argv)
    check_window = LabelVideoWindow(root_folder, video_path)
    check_window.show()
    sys.exit(app.exec_())


if __name__ == "__main__":
    main()
